#include <iostream>
#include <vector>

#include <map>
#include <set>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include "cluster.h"

#include "random.h"

using namespace std;
using namespace sds;

int state = 50;

double maxDistance = 128.0;

Random rng;

void renumber_dfs(int i, int j, int n, int c, 
		  SDataView<unsigned char>& col,
		  SDataView<unsigned int>& num) {
  if (i < 0 || j < 0 || 
      i >= col.getSize()[0] || j >= col.getSize()[1] ||
      col[i][j] != c) return;
  num[i][j] = n;
  col[i][j] = 255;
  renumber_dfs(i - 1, j, n, c, col, num);
  renumber_dfs(i + 1, j, n, c, col, num);
  renumber_dfs(i, j - 1, n, c, col, num);
  renumber_dfs(i, j + 1, n, c, col, num);
}

void renumber(SDataView<unsigned char>& col,
	      SDataView<unsigned int>& num, int& n) {
  n = 1;
  for (int i = 0; i < col.getSize()[0]; ++i) {
    for (int j = 0; j < col.getSize()[1]; ++j) {
      if (col[i][j] == 255) continue;
      if (col[i][j] != 0) {
	renumber_dfs(i, j, n, col[i][j], col, num);
	++n;
      } else {
	num[i][j] = 0;
      }
    }
  }
}

void Read(const string& filename, const std::string& segname,
	  SDataView<Point>& src, SDataView<unsigned int>& seg,
	  vector<SInteger>& size, int& segcnt) {

  SFile segfile(segname);
  SDataView<unsigned char> segcol = 
    segfile.getDataset(0).getDataView<unsigned char>();
  size = segcol.getSize();
  
  seg = SDataView<unsigned int>(size);

  renumber(segcol, seg, segcnt);

  SDataView<unsigned char> source[CH];
  SFile file(filename);
  for (int k = 0; k < CH; ++k) {
    SDataset dataset = file.getDataset(k);
    source[k] = dataset.getDataView<unsigned char>();
  }

  src = SDataView<Point>(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      src[i][j] = point;
    }
  }  
}

void Write(const string& filename,
	   SDataView<Point>& src,
	   const vector<Cluster>& sta,
	   const vector<SInteger>& size) {

  SFile file(filename, SFile::CREATE);

  SDataView<unsigned char> targetView(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point item = src[i][j];
      int minPos = 0;
      double minVal = distance_rp(item, sta[0].center());
      for (int l = 1; l < state; ++l) {
	double val = distance_rp(item, sta[l].center());
	if (val < minVal) {
	  minVal = val;
	  minPos = l;
	}
      }
      if (minVal < maxDistance) {
	targetView[i][j] = (unsigned char)minPos + 1;
      } else {
	targetView[i][j] = 0;
      }
    }
  }

  SDataset dataset = file.createDataset("clustered", UINT8, size);
  
  dataset.setDataView<unsigned char>(targetView);
}


int main(int argc, const char *argv[]) {
  assert(argc == 4);

  SDataView<unsigned int> seg;
  SDataView<Point> src;
  vector<SInteger> size;

  int n;

  Read(argv[1], argv[2], src, seg, size, n);
  
  --n;

  vector<unsigned char> clusters(n);
  vector<Cluster> segments(n);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (seg[i][j] != 0) {
	segments[seg[i][j] - 1] += src[i][j];
      }
    }
  }

  vector<Cluster> states(state);

  for (int i = 0; i < n; ++i) {
    char cl = (char)(state * rng.uniform()); 
    clusters[i] = cl + 1;
    states[(int)cl] += segments[i];
  }

  const int length = 60;

  for (int k = 0; k < length; ++k) {
    cout << "Step: " << k << endl;
    double t = 3.0 / log (1.0 + k * 0.08);
    for (int i = 0; i < n; ++i) {
      if (clusters[i] != 0) {
	states[clusters[i] - 1] -= segments[i];
      }
      double probs[state];
      double sum = 0.0;
      for (int l = 0; l < (int)state; ++l) {
	double dist = 
	  states[l].size() != 0 ? distance_avg(states[l], segments[i]) : 0.0;
	probs[l] = exp( - dist / t);
	sum += probs[l];
      }
      sum += exp( - maxDistance / t);
      double prob = rng.uniform() * sum;
      clusters[i] = 0;
      for (int l = 0; l < (int)state; ++l) {
	if (prob < probs[l]) {
	  clusters[i] = l + 1;
	  states[l] += segments[i];
	  break;
	}
	prob -= probs[l];
      }
    }
  }

  //  Write(argv[3], seg, clusters, size);
  Write(argv[3], src, states, size);
  
  return 0;
}
